"""图论与结构矩阵模块
包含邻接矩阵、拉普拉斯矩阵、图的谱分析、图算法等
支持有向图、无向图、加权图的各种矩阵表示和分析
"""
import numpy as np
from scipy import sparse
from scipy.sparse.linalg import eigsh, eigs
from typing import Union, List, Tuple, Optional, Dict
import warnings

class GraphTheory:
    """图论矩阵分析类"""
    
    # ==================== 图的基本矩阵表示 ====================
    
    @staticmethod
    def create_adjacency_matrix(edges: List[Tuple[int, int]], n_nodes: int, 
                               weights: Optional[List[float]] = None, 
                               directed: bool = False) -> np.ndarray:
        """
        构建邻接矩阵
        :param edges: 边列表 [(i, j), ...]
        :param n_nodes: 节点数量
        :param weights: 边权重（可选）
        :param directed: 是否为有向图
        :return: 邻接矩阵
        """
        A = np.zeros((n_nodes, n_nodes))
        
        for idx, (i, j) in enumerate(edges):
            weight = weights[idx] if weights else 1.0
            A[i, j] = weight
            if not directed:
                A[j, i] = weight
        
        return A
    
    @staticmethod
    def create_sparse_adjacency_matrix(edges: List[Tuple[int, int]], n_nodes: int,
                                      weights: Optional[List[float]] = None,
                                      directed: bool = False) -> sparse.csr_matrix:
        """
        构建稀疏邻接矩阵
        :param edges: 边列表
        :param n_nodes: 节点数量
        :param weights: 边权重（可选）
        :param directed: 是否为有向图
        :return: 稀疏邻接矩阵
        """
        row_indices = []
        col_indices = []
        data = []
        
        for idx, (i, j) in enumerate(edges):
            weight = weights[idx] if weights else 1.0
            row_indices.append(i)
            col_indices.append(j)
            data.append(weight)
            
            if not directed:
                row_indices.append(j)
                col_indices.append(i)
                data.append(weight)
        
        return sparse.csr_matrix((data, (row_indices, col_indices)), shape=(n_nodes, n_nodes))
    
    @staticmethod
    def degree_matrix(adjacency: Union[np.ndarray, sparse.spmatrix]) -> Union[np.ndarray, sparse.spmatrix]:
        """
        计算度矩阵
        :param adjacency: 邻接矩阵
        :return: 度矩阵
        """
        if sparse.issparse(adjacency):
            degrees = np.array(adjacency.sum(axis=1)).flatten()
            return sparse.diags(degrees, format='csr')
        else:
            degrees = np.sum(adjacency, axis=1)
            return np.diag(degrees)
    
    @staticmethod
    def incidence_matrix(edges: List[Tuple[int, int]], n_nodes: int, 
                        directed: bool = False) -> np.ndarray:
        """
        构建关联矩阵
        :param edges: 边列表
        :param n_nodes: 节点数量
        :param directed: 是否为有向图
        :return: 关联矩阵
        """
        n_edges = len(edges)
        B = np.zeros((n_nodes, n_edges))
        
        for edge_idx, (i, j) in enumerate(edges):
            if directed:
                B[i, edge_idx] = -1  # 出边
                B[j, edge_idx] = 1   # 入边
            else:
                B[i, edge_idx] = 1
                B[j, edge_idx] = 1
        
        return B
    
    # ==================== 拉普拉斯矩阵 ====================
    
    @staticmethod
    def laplacian_matrix(adjacency: Union[np.ndarray, sparse.spmatrix]) -> Union[np.ndarray, sparse.spmatrix]:
        """
        计算拉普拉斯矩阵 L = D - A
        :param adjacency: 邻接矩阵
        :return: 拉普拉斯矩阵
        """
        degree = GraphTheory.degree_matrix(adjacency)
        return degree - adjacency
    
    @staticmethod
    def normalized_laplacian(adjacency: Union[np.ndarray, sparse.spmatrix], 
                           method: str = 'symmetric') -> Union[np.ndarray, sparse.spmatrix]:
        """
        计算归一化拉普拉斯矩阵
        :param adjacency: 邻接矩阵
        :param method: 归一化方法 ('symmetric', 'random_walk')
        :return: 归一化拉普拉斯矩阵
        """
        if sparse.issparse(adjacency):
            degrees = np.array(adjacency.sum(axis=1)).flatten()
            degrees[degrees == 0] = 1  # 避免除零
            
            if method == 'symmetric':
                # L_sym = D^(-1/2) L D^(-1/2)
                degree_sqrt_inv = sparse.diags(1.0 / np.sqrt(degrees), format='csr')
                laplacian = GraphTheory.laplacian_matrix(adjacency)
                return degree_sqrt_inv @ laplacian @ degree_sqrt_inv
            elif method == 'random_walk':
                # L_rw = D^(-1) L
                degree_inv = sparse.diags(1.0 / degrees, format='csr')
                laplacian = GraphTheory.laplacian_matrix(adjacency)
                return degree_inv @ laplacian
        else:
            degrees = np.sum(adjacency, axis=1)
            degrees[degrees == 0] = 1  # 避免除零
            
            if method == 'symmetric':
                degree_sqrt_inv = np.diag(1.0 / np.sqrt(degrees))
                laplacian = GraphTheory.laplacian_matrix(adjacency)
                return degree_sqrt_inv @ laplacian @ degree_sqrt_inv
            elif method == 'random_walk':
                degree_inv = np.diag(1.0 / degrees)
                laplacian = GraphTheory.laplacian_matrix(adjacency)
                return degree_inv @ laplacian
        
        raise ValueError(f"Unknown normalization method: {method}")
    
    @staticmethod
    def signless_laplacian(adjacency: Union[np.ndarray, sparse.spmatrix]) -> Union[np.ndarray, sparse.spmatrix]:
        """
        计算无符号拉普拉斯矩阵 Q = D + A
        :param adjacency: 邻接矩阵
        :return: 无符号拉普拉斯矩阵
        """
        degree = GraphTheory.degree_matrix(adjacency)
        return degree + adjacency
    
    # ==================== 图的谱分析 ====================
    
    @staticmethod
    def graph_spectrum(matrix: Union[np.ndarray, sparse.spmatrix], 
                      k: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        计算图矩阵的谱（特征值和特征向量）
        :param matrix: 图矩阵
        :param k: 计算前k个特征值（对稀疏矩阵有效）
        :return: (特征值, 特征向量)
        """
        if sparse.issparse(matrix):
            if k is None:
                k = min(matrix.shape[0] - 1, 10)
            eigenvals, eigenvecs = eigsh(matrix, k=k, which='SM')
            return eigenvals, eigenvecs
        else:
            eigenvals, eigenvecs = np.linalg.eigh(matrix)
            if k is not None:
                return eigenvals[:k], eigenvecs[:, :k]
            return eigenvals, eigenvecs
    
    @staticmethod
    def algebraic_connectivity(laplacian: Union[np.ndarray, sparse.spmatrix]) -> float:
        """
        计算代数连通度（第二小特征值）
        :param laplacian: 拉普拉斯矩阵
        :return: 代数连通度
        """
        if sparse.issparse(laplacian):
            eigenvals, _ = eigsh(laplacian, k=2, which='SM')
            return eigenvals[1]
        else:
            eigenvals = np.linalg.eigvals(laplacian)
            eigenvals = np.sort(eigenvals)
            return eigenvals[1]
    
    @staticmethod
    def fiedler_vector(laplacian: Union[np.ndarray, sparse.spmatrix]) -> np.ndarray:
        """
        计算Fiedler向量（第二小特征值对应的特征向量）
        :param laplacian: 拉普拉斯矩阵
        :return: Fiedler向量
        """
        if sparse.issparse(laplacian):
            eigenvals, eigenvecs = eigsh(laplacian, k=2, which='SM')
            return eigenvecs[:, 1]
        else:
            eigenvals, eigenvecs = np.linalg.eigh(laplacian)
            return eigenvecs[:, 1]
    
    @staticmethod
    def graph_eigenmaps(laplacian: Union[np.ndarray, sparse.spmatrix], 
                       k: int = 2) -> np.ndarray:
        """
        图的特征映射（谱嵌入）
        :param laplacian: 拉普拉斯矩阵
        :param k: 嵌入维度
        :return: 特征向量矩阵
        """
        if sparse.issparse(laplacian):
            eigenvals, eigenvecs = eigsh(laplacian, k=k+1, which='SM')
            return eigenvecs[:, 1:k+1]  # 跳过第一个特征向量（常向量）
        else:
            eigenvals, eigenvecs = np.linalg.eigh(laplacian)
            return eigenvecs[:, 1:k+1]
    
    # ==================== 图的连通性分析 ====================
    
    @staticmethod
    def connected_components(adjacency: Union[np.ndarray, sparse.spmatrix]) -> Tuple[int, np.ndarray]:
        """
        计算连通分量
        :param adjacency: 邻接矩阵
        :return: (连通分量数, 节点标签)
        """
        from scipy.sparse.csgraph import connected_components
        
        if not sparse.issparse(adjacency):
            adjacency = sparse.csr_matrix(adjacency)
        
        n_components, labels = connected_components(adjacency, directed=False)
        return n_components, labels
    
    @staticmethod
    def is_connected(adjacency: Union[np.ndarray, sparse.spmatrix]) -> bool:
        """
        判断图是否连通
        :param adjacency: 邻接矩阵
        :return: 是否连通
        """
        n_components, _ = GraphTheory.connected_components(adjacency)
        return n_components == 1
    
    @staticmethod
    def graph_diameter(adjacency: Union[np.ndarray, sparse.spmatrix]) -> float:
        """
        计算图的直径（最短路径的最大值）
        :param adjacency: 邻接矩阵
        :return: 图的直径
        """
        from scipy.sparse.csgraph import shortest_path
        
        if not sparse.issparse(adjacency):
            adjacency = sparse.csr_matrix(adjacency)
        
        distances = shortest_path(adjacency, directed=False)
        finite_distances = distances[np.isfinite(distances)]
        return np.max(finite_distances) if len(finite_distances) > 0 else np.inf
    
    # ==================== 图的中心性度量 ====================
    
    @staticmethod
    def degree_centrality(adjacency: Union[np.ndarray, sparse.spmatrix]) -> np.ndarray:
        """
        计算度中心性
        :param adjacency: 邻接矩阵
        :return: 度中心性向量
        """
        if sparse.issparse(adjacency):
            degrees = np.array(adjacency.sum(axis=1)).flatten()
        else:
            degrees = np.sum(adjacency, axis=1)
        
        n = adjacency.shape[0]
        return degrees / (n - 1) if n > 1 else degrees
    
    @staticmethod
    def eigenvector_centrality(adjacency: Union[np.ndarray, sparse.spmatrix], 
                              max_iter: int = 1000, tol: float = 1e-6) -> np.ndarray:
        """
        计算特征向量中心性
        :param adjacency: 邻接矩阵
        :param max_iter: 最大迭代次数
        :param tol: 收敛容差
        :return: 特征向量中心性
        """
        if sparse.issparse(adjacency):
            eigenvals, eigenvecs = eigs(adjacency, k=1, which='LR', maxiter=max_iter, tol=tol)
            centrality = np.abs(eigenvecs[:, 0])
        else:
            eigenvals, eigenvecs = np.linalg.eig(adjacency)
            max_eigenval_idx = np.argmax(eigenvals.real)
            centrality = np.abs(eigenvecs[:, max_eigenval_idx])
        
        return centrality / np.sum(centrality)
    
    @staticmethod
    def pagerank_centrality(adjacency: Union[np.ndarray, sparse.spmatrix], 
                           damping: float = 0.85, max_iter: int = 1000, 
                           tol: float = 1e-6) -> np.ndarray:
        """
        计算PageRank中心性
        :param adjacency: 邻接矩阵
        :param damping: 阻尼因子
        :param max_iter: 最大迭代次数
        :param tol: 收敛容差
        :return: PageRank中心性
        """
        n = adjacency.shape[0]
        
        # 构建转移矩阵
        if sparse.issparse(adjacency):
            out_degrees = np.array(adjacency.sum(axis=1)).flatten()
            out_degrees[out_degrees == 0] = 1  # 避免除零
            transition = adjacency.T / out_degrees
        else:
            out_degrees = np.sum(adjacency, axis=1)
            out_degrees[out_degrees == 0] = 1
            transition = adjacency.T / out_degrees[:, np.newaxis]
        
        # 幂迭代法
        pagerank = np.ones(n) / n
        
        for _ in range(max_iter):
            if sparse.issparse(transition):
                new_pagerank = damping * (transition @ pagerank) + (1 - damping) / n
            else:
                new_pagerank = damping * (transition @ pagerank) + (1 - damping) / n
            
            if np.linalg.norm(new_pagerank - pagerank) < tol:
                break
            
            pagerank = new_pagerank
        
        return pagerank
    
    # ==================== 图的聚类 ====================
    
    @staticmethod
    def spectral_clustering(adjacency: Union[np.ndarray, sparse.spmatrix], 
                           n_clusters: int, method: str = 'normalized') -> np.ndarray:
        """
        谱聚类
        :param adjacency: 邻接矩阵
        :param n_clusters: 聚类数
        :param method: 拉普拉斯矩阵类型 ('unnormalized', 'normalized', 'random_walk')
        :return: 聚类标签
        """
        from sklearn.cluster import KMeans
        
        if method == 'unnormalized':
            laplacian = GraphTheory.laplacian_matrix(adjacency)
        elif method == 'normalized':
            laplacian = GraphTheory.normalized_laplacian(adjacency, method='symmetric')
        elif method == 'random_walk':
            laplacian = GraphTheory.normalized_laplacian(adjacency, method='random_walk')
        else:
            raise ValueError(f"Unknown method: {method}")
        
        # 计算特征向量
        if sparse.issparse(laplacian):
            eigenvals, eigenvecs = eigsh(laplacian, k=n_clusters, which='SM')
        else:
            eigenvals, eigenvecs = np.linalg.eigh(laplacian)
            eigenvecs = eigenvecs[:, :n_clusters]
        
        # K-means聚类
        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        labels = kmeans.fit_predict(eigenvecs)
        
        return labels
    
    @staticmethod
    def modularity_matrix(adjacency: Union[np.ndarray, sparse.spmatrix]) -> Union[np.ndarray, sparse.spmatrix]:
        """
        计算模块度矩阵
        :param adjacency: 邻接矩阵
        :return: 模块度矩阵
        """
        if sparse.issparse(adjacency):
            degrees = np.array(adjacency.sum(axis=1)).flatten()
            m = adjacency.sum() / 2  # 边数
            expected = sparse.csr_matrix(np.outer(degrees, degrees) / (2 * m))
            return adjacency - expected
        else:
            degrees = np.sum(adjacency, axis=1)
            m = np.sum(adjacency) / 2
            expected = np.outer(degrees, degrees) / (2 * m)
            return adjacency - expected
    
    # ==================== 图的可视化辅助 ====================
    
    @staticmethod
    def graph_layout_spring(adjacency: Union[np.ndarray, sparse.spmatrix], 
                           dim: int = 2, iterations: int = 50) -> np.ndarray:
        """
        弹簧布局算法
        :param adjacency: 邻接矩阵
        :param dim: 布局维度
        :param iterations: 迭代次数
        :return: 节点坐标
        """
        n = adjacency.shape[0]
        pos = np.random.rand(n, dim)
        
        for _ in range(iterations):
            # 计算斥力
            for i in range(n):
                for j in range(i + 1, n):
                    diff = pos[i] - pos[j]
                    dist = np.linalg.norm(diff)
                    if dist > 0:
                        force = diff / (dist ** 3)
                        pos[i] += force * 0.1
                        pos[j] -= force * 0.1
            
            # 计算引力
            if sparse.issparse(adjacency):
                rows, cols = adjacency.nonzero()
                for i, j in zip(rows, cols):
                    if i < j:
                        diff = pos[i] - pos[j]
                        dist = np.linalg.norm(diff)
                        if dist > 0:
                            force = diff * dist * 0.01
                            pos[i] -= force
                            pos[j] += force
            else:
                for i in range(n):
                    for j in range(i + 1, n):
                        if adjacency[i, j] > 0:
                            diff = pos[i] - pos[j]
                            dist = np.linalg.norm(diff)
                            if dist > 0:
                                force = diff * dist * 0.01
                                pos[i] -= force
                                pos[j] += force
        
        return pos
    
    # ==================== 图的统计信息 ====================
    
    @staticmethod
    def graph_statistics(adjacency: Union[np.ndarray, sparse.spmatrix]) -> Dict:
        """
        计算图的统计信息
        :param adjacency: 邻接矩阵
        :return: 统计信息字典
        """
        n_nodes = adjacency.shape[0]
        
        if sparse.issparse(adjacency):
            n_edges = adjacency.nnz // 2  # 假设无向图
            degrees = np.array(adjacency.sum(axis=1)).flatten()
        else:
            n_edges = int(np.sum(adjacency) // 2)
            degrees = np.sum(adjacency, axis=1)
        
        density = 2 * n_edges / (n_nodes * (n_nodes - 1)) if n_nodes > 1 else 0
        
        return {
            'n_nodes': n_nodes,
            'n_edges': n_edges,
            'density': density,
            'avg_degree': np.mean(degrees),
            'max_degree': np.max(degrees),
            'min_degree': np.min(degrees),
            'degree_std': np.std(degrees),
            'is_connected': GraphTheory.is_connected(adjacency),
            'n_components': GraphTheory.connected_components(adjacency)[0]
        }